//
//  RBTree.swift
//  others
//
//  Created by lhc on 2024/8/27.
//

import Foundation

/// 红黑树
class RBTree<T: Comparable, E>: BalanceBinaryTree<T,E> {
    
    /// 右旋转
    ///
    /// - node结点有一个左红结点
    /// - 左红结点又有一个新的左红结点
    /// - 返回的node是第一个左红结点，旧node变成新node的右红结点
    private func rightRotate(node: Node?) -> Node? {
        let x = node?.left
        node?.left = x?.right
        x?.right = node
        
        x?.color = node?.color ?? .black
        node?.color = .red
        return x
    }
    
    /// 左旋转
    ///
    /// - node结点添加一个大于key的新结点后需要左旋转
    /// - 返回的node是新添加的结点，旧node变成新node的左红结点
    private func leftRotate(node: Node?) -> Node? {
        let x = node?.right
        node?.right = x?.left
        x?.left = node
        
        // 旋转后x为根，他替代之前的node，颜色为之前的node的颜色
        x?.color = node?.color ?? .black
        // node变成左红结点
        node?.color = .red
        return x
    }
    
    /// 三个结点把中间的变红上升， 左右子树变成黑结点
    private func flipColor(node: Node?) {
        node?.color = .red
        node?.left?.color = .black
        node?.right?.color = .black
    }
    
    /// 判断node是否是红
    private func isRed(node: Node?) -> Bool {
        if let color = node?.color {
            return color == .red
        }else {
            return false
        }
    }
    
    func add(key: T, value: E?) {
        self.root = self.add(key: key, value: value, node: self.root)
    }
    
    private func add(key: T, value: E?, node: Node?) -> Node? {
        guard let node = node else {
            self.size += 1
            return Node(key: key, value: value)
        }
        
        
        if key < node.key {
            node.left = self.add(key: key, value: value, node: node.left)
        }else if key > node.key {
            node.right = self.add(key: key, value: value, node: node.right)
        }else {
            // key == node.key
            node.value = value
        }
        
        var resultNode: Node? = nil
        // 如果右红左黑
        if self.isRed(node: node.right) && !self.isRed(node: node.left){
            resultNode = self.leftRotate(node: node)
        }
        // 如果左红 左红的左依然红
        if self.isRed(node: node.left) && self.isRed(node: node.left?.left) {
            resultNode = self.rightRotate(node: node)
        }
        
        // 如果左右都红
        if self.isRed(node: node.left) && self.isRed(node: node.right) {
            self.flipColor(node: resultNode)
        }
        
        // 更新node高度
        resultNode?.height = max(self.getHeightWithNode(resultNode?.left), self.getHeightWithNode(resultNode?.right)) + 1
        
        return resultNode
        
    }
}
